use serde::{Deserialize, Serialize};
use serde_with::serde_as;

use aspect_codec::as_with::AsCodecHex;
use aspect_codec::prelude::*;

use super::crypto;
use super::crypto::SigPurpose;
use super::crypto::IDSIG_PURPOSE_ATTESTATION;
use super::crypto::IDSIG_PURPOSE_PERSONAL;
use super::device::*;
use super::errors;

const TRUST_LIFETIME: u64 = 1000 * 60 * 60 * 24 * 30; // 30 days
const NUM_PREKEYS: usize = 8;
const ANN_SM_LEN: usize = 64; // 32 (ident) + 32 (ann hash)
const UPD_SM_LEN: usize = 68; // 32 (ident) + 4 (seqno) + 32 (upd hash)

/// Used to prepare an identity as it's created.  Eventually will help manage
/// the state involved in getting identity sponsorships.
pub struct SignerBuilder {
    ident_kp: crypto::IdentKeypair,
    inbox_kp: crypto::InboxKeypair,

    short_name: String,
}

impl SignerBuilder {
    pub fn new(short_name: String) -> Self {
        Self {
            ident_kp: crypto::gen_ident_keypair(crypto::KeyType::Ec),
            inbox_kp: crypto::gen_inbox_keypair(crypto::KeyType::Ec),
            short_name,
        }
    }

    /// Finalizes the state of the builder into a real signer and a DIR package
    /// containing an announcement.
    pub fn finalize(self) -> Result<(IdentSigner, DirPackage), errors::Error> {
        let now = now_ts();
        let ident_pk = self.ident_kp.to_pub();
        let inbox_pk = self.inbox_kp.to_pub();
        let id = DeviceIdent::from_ident_pubkey(&ident_pk);

        let adata = AnnounceData::new(now, ident_pk, self.short_name, None);
        let u0data = UpdateData::new(now, TRUST_LIFETIME, inbox_pk, String::new(), Vec::new());

        // Now sign the packages.
        let apkg = create_announce_package(&self.ident_kp, &id, &adata);
        let u0pkg = create_update_package(&self.ident_kp, &id, 0, &u0data);
        let full_pkg = DirPackage::new(apkg, u0pkg);

        let signer = IdentSigner {
            ident_kp: self.ident_kp,
            inbox_kp: self.inbox_kp,
            last_package: full_pkg.clone(),
        };

        Ok((signer, full_pkg))
    }
}

/// Signer used to produce new device update packages.
#[serde_as]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct IdentSigner {
    #[serde_as(as = "AsCodecHex")]
    ident_kp: crypto::IdentKeypair,

    #[serde_as(as = "AsCodecHex")]
    inbox_kp: crypto::InboxKeypair,

    #[serde_as(as = "AsCodecHex")]
    last_package: DirPackage,
}

impl IdentSigner {
    pub fn cur_package(&self) -> &DirPackage {
        &self.last_package
    }

    pub fn announcement(&self) -> &AnnPackage {
        self.cur_package().announcement()
    }

    pub fn last_update(&self) -> &UpdatePackage {
        self.cur_package().last_update()
    }

    /// Returns the current sequence number.  Ie. the seqno on the last published update.
    pub fn seqno(&self) -> u32 {
        self.last_update().seqno()
    }

    fn get_device_ident(&self) -> DeviceIdent {
        let pk = self.ident_kp.to_pub();
        DeviceIdent::from_ident_pubkey(&pk)
    }

    /// Creates an new update, incrementing the seqno and internal state.  The
    /// resulting package should only be published *after* the updated signer
    /// state has been committed to disk.  This is crash-safe as long as we
    /// re-announce our currently saved dirpackage to peers every time we start.
    ///
    /// Currently only allows changing the update comment.
    pub fn create_update(&mut self, comment: String) -> DirPackage {
        let ts = now_ts();

        let ann = self.last_package.announcement().clone();
        let seqno = self.seqno() + 1;
        let id = self.get_device_ident();

        let udata = UpdateData::new(
            ts,
            TRUST_LIFETIME,
            self.inbox_kp.to_pub(),
            comment,
            Vec::new(),
        );

        let upkg = create_update_package(&self.ident_kp, &id, seqno, &udata);
        let pkg = DirPackage::new(ann, upkg);
        self.last_package = pkg.clone();
        pkg
    }

    /// Creates a signature for a general message using the identity signing
    /// key.
    pub fn sign_generic_message(&self, purpose: SigPurpose, msg: &[u8]) -> crypto::Signature {
        crypto::sign_ident_msg(&self.ident_kp, purpose, msg)
    }

    /// Creates a serialized personal message in the appropriate container.
    pub fn create_personal_message(&self, msg: String) -> SignedMessage {
        let sig = self.sign_generic_message(IDSIG_PURPOSE_PERSONAL, msg.as_bytes());
        SignedMessage {
            sig,
            message: msg.into_bytes(),
        }
    }

    /// Creates a device attestation message for some challenge.
    pub fn create_device_attestation(&self, challenge: &[u8]) -> DeviceAttestation {
        let sig = self.sign_generic_message(IDSIG_PURPOSE_ATTESTATION, challenge);
        DeviceAttestation::new(self.announcement().clone(), sig)
    }
}

/// Signs a package from the announcement and stuff.  Avoids recomputing the DI.
fn create_announce_package(
    kp: &crypto::IdentKeypair,
    ident: &DeviceIdent,
    data: &AnnounceData,
) -> AnnPackage {
    let ann_enc = aspect_codec::util::encode_to_vec(data).expect("signer: enc ann");
    let pkg_msg = gen_announce_sig_msg(ident, &ann_enc);
    let sig = crypto::sign_ident_msg(kp, crypto::IDSIG_PURPOSE_ANN, &pkg_msg);
    AnnPackage::new(sig, ann_enc)
}

/// Signs a package from the update data and seqno and stuff.  Avoids recomputing the DI.
fn create_update_package(
    kp: &crypto::IdentKeypair,
    ident: &DeviceIdent,
    seqno: u32,
    data: &UpdateData,
) -> UpdatePackage {
    let upd_enc = aspect_codec::util::encode_to_vec(data).expect("signer: enc update");
    let pkg_msg = gen_update_sig_msg(ident, seqno, &upd_enc);
    let sig = crypto::sign_ident_msg(kp, crypto::IDSIG_PURPOSE_IDENT, &pkg_msg);
    UpdatePackage::new(*ident, seqno, sig, upd_enc)
}

/// Verifies that an announcement package is internally consistent.
pub fn verify_ann_package_integrity(ann: &AnnPackage) -> bool {
    let Ok(data) = ann.try_decode() else {
        return false;
    };

    let ident = data.get_ident();
    let ann_msg = gen_announce_sig_msg(&ident, ann.announcement_raw());
    if !crypto::verify_ident_msg(
        ann.sig(),
        data.ident_pk(),
        crypto::IDSIG_PURPOSE_ANN,
        &ann_msg,
    ) {
        return false;
    }

    true
}

/// Verifies if a DIR package matches the provided pubkey and is correctly signed.
pub fn verify_dir_package(pk: &crypto::IdentPubkey, dir: &UpdatePackage) -> bool {
    // Check ident matches.  If this didn't match then things would be really
    // weird and we wouldn't know what to trust.
    if DeviceIdent::from_ident_pubkey(pk) != *dir.ident() {
        return false;
    }

    // Check signature.  The whole point of DIs is that the signatures keep
    // everything trustworthy.
    let pkg_msg = gen_update_sig_msg(dir.ident(), dir.seqno(), dir.record_raw());
    if !crypto::verify_ident_msg(dir.sig(), pk, crypto::IDSIG_PURPOSE_IDENT, &pkg_msg) {
        return false;
    }

    true
}

/// Verifies that a signed message was appropriately signed by the specified pubkey.
pub fn verify_personal_signed_message(ident: &crypto::IdentPubkey, smsg: &SignedMessage) -> bool {
    crypto::verify_ident_msg(smsg.sig(), ident, IDSIG_PURPOSE_PERSONAL, smsg.message())
}

/// Generates an announcement signature message.  It's a flat byte buffer.
fn gen_announce_sig_msg(ident: &DeviceIdent, ann_data: &[u8]) -> [u8; ANN_SM_LEN] {
    let h = aspect_codec::util::compute_sha256_raw(ann_data);
    let mut pkg_msg = [0; ANN_SM_LEN];
    pkg_msg[..32].copy_from_slice(&ident.as_ref()[..]);
    pkg_msg[32..].copy_from_slice(&h);
    pkg_msg
}

/// Generates an update signature message.  It's a flat byte buffer.
fn gen_update_sig_msg(ident: &DeviceIdent, seqno: u32, upd_record: &[u8]) -> [u8; UPD_SM_LEN] {
    let h = aspect_codec::util::compute_sha256_raw(upd_record);
    let mut pkg_msg = [0; UPD_SM_LEN];
    pkg_msg[..32].copy_from_slice(&ident.as_ref()[..]);
    pkg_msg[32..36].copy_from_slice(&seqno.to_be_bytes()[..]);
    pkg_msg[36..68].copy_from_slice(&h);
    pkg_msg
}

fn now_ts() -> u64 {
    std::time::UNIX_EPOCH.elapsed().unwrap().as_millis() as u64
}

/// Generic signed personal message container.  Does not identify itself.
#[derive(Clone, Debug)]
pub struct SignedMessage {
    sig: crypto::Signature,
    message: Vec<u8>,
}

impl SignedMessage {
    pub fn sig(&self) -> &crypto::Signature {
        &self.sig
    }

    pub fn message(&self) -> &[u8] {
        &self.message
    }

    pub fn into_message(self) -> Vec<u8> {
        self.message
    }
}

aspect_codec::inst_struct_codec! {
    SignedMessage {
        sig: crypto::Signature,
        message: Vec<u8>,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_signer_basic() {
        let builder = SignerBuilder::new("testuser".to_owned());
        let (mut signer, dir0) = builder.finalize().expect("test: init signer");

        assert_eq!(
            aspect_codec::util::encode_to_vec(signer.announcement()).unwrap(),
            aspect_codec::util::encode_to_vec(dir0.announcement()).unwrap()
        );

        assert_eq!(signer.seqno(), 0);
        assert_eq!(dir0.last_update().seqno(), 0);

        let dir1 = signer.create_update("hello".to_owned());

        assert_eq!(
            aspect_codec::util::encode_to_vec(signer.announcement()).unwrap(),
            aspect_codec::util::encode_to_vec(dir1.announcement()).unwrap()
        );

        assert_eq!(signer.seqno(), 1);
        assert_eq!(dir1.last_update().seqno(), 1);

        let dir2 = signer.create_update("hello again".to_owned());

        assert_eq!(
            aspect_codec::util::encode_to_vec(signer.announcement()).unwrap(),
            aspect_codec::util::encode_to_vec(dir2.announcement()).unwrap()
        );

        assert_eq!(signer.seqno(), 2);
        assert_eq!(dir2.last_update().seqno(), 2);

        eprintln!("{:?}", dir2);
    }

    #[test]
    fn test_dev_attestation() {
        let builder = SignerBuilder::new("testuser".to_owned());
        let (signer, _dir0) = builder.finalize().expect("test: init signer");

        let attest_msg = &[0, 1, 2, 3, 4, 5, 6];
        let attestation = signer.create_device_attestation(attest_msg);

        assert!(verify_attestation(&attestation, attest_msg));
    }
}
